#!/usr/bin/env python
import os
import sys
import re
import argparse
import shutil
class crysInpGen():
    def __init__(self,title):
        self.title=title
        self.cell=''
        self.coord=[]
        self.geom=''
        self.shrink=''
        self.basis=''
        self.scf=''
        self.demension=1
        self.symGroup=1
        self.opt=''
        self.dft=''
        self.atoms=0
        self.guess=''
        self.cycle=''
        self.supercell=''
        self.grimme=''
        self.elements=dict(H=1,He=2,Li=3,Be=4,B=5,C=6,N=7,O=8,F=9,Na=11,Mg=12,Al=13,Si=14,P=15,S=16,Cl=17,Br=35)
        self.gd2Param=dict(pbepbe='0.75',b3lyp='1.05',blyp='1.2',beckep86='1.05',beckelyp='1.2')
        self.elementGD2={1:'0.14 1.001',2:'0.08 1.012',3:'1.61 0.825',4:'1.61 1.408',5:'3.13 1.485',6:'1.75 1.452',7:'1.23 1.397',8:'0.70 1,342',9:'0.75 1.287', \
        10:'0.63 1.243',11:'5.71 1.144',12:'5.71 1.364',13:'10.79 1.639',14:'9.23 1.716',15:'7.84 1.705',16:'5.57 1.683',17:'5.07 1.639',18:'4.61 1.595',35:'12.47 1.749'}
        self.D2sym=dict()
    
    def readcif(self,cifFileName):
        cifFile=open(cifFileName,'r')  
        # read coordinate and cell parameters from cif file         
        for line in cifFile:
            if re.search('_cell_length_a',line):
                self.cell_a=line.split()[1]
            if re.search('_cell_length_b',line):
                self.cell_b=line.split()[1]
            if re.search('_cell_length_c',line):
                self.cell_c=line.split()[1]
            if re.search('_cell_angle_alpha',line):
                self.cell_alpha=line.split()[1]       
            if re.search('_cell_angle_beta',line):
                self.cell_beta=line.split()[1]   
            if re.search('_cell_angle_gamma',line):
                self.cell_gamma=line.split()[1] 
            if '_symmetry_Int_Tables_number' in line:
                self.ITN=line.split()[1]              
            if re.search('-?\d+\.\d+\s+-?\d+\.\d+\s+-?\d+.\d+\s?',line):
                self.atoms=self.atoms+1
                coord=[self.elements[line.split()[1]],float(line.split()[2]),float(line.split()[3]),float(line.split()[4])]
                self.coord.append(coord)
        cifFile.close()

    def readgui(self,guiFileName):
        guiFile=open(guiFileName,'r')
        for line in guiFile:
            if re.search('\s+\d+\s+-?\d+\.\d+\s+-?\d+\.\d+\s+-?\d+\.\d+',line):
                coord=line.split()
                self.coord.append(coord)

    
    def crysParamParser(self,args,external):
        if len(args.dft.split('|'))==2:
            exchange=dft.split('|')[0]
            correlat=dft.split('|')[1]
            self.dft="DFT\nEXCHANGE\n"+exchange+"\nCORRELAT\n"+correlat+"\nEND\n"
        if len(args.dft.split('|'))==1:                
            self.dft="DFT\n"+dft+"\nEND\n"
    #parser grimme gd2 input
        if args.grimme:
            dftName=''.join(dft.split('|')).lower()
            elements=[]
            for coord in self.coord:
                if int(coord[0]) not in elements:
                    elements.append(int(coord[0]))
            if not elements and args.grimme.lower() != 'a':
                elements=map(int,args.grimme.split(','))
            elif not elements and args.grimme.lower() == 'a':
                print 'Warning!!! no element found when trying to determine their grimme D2 parameter'
            elementParam=''
            for element in elements:
                elementParam=elementParam+str(element)+' '+self.elementGD2[element]+'\n'
            dftParam=''
            if dftName in self.gd2Param:
                dftParam=self.gd2Param[dftName]+' 20. 25.\n'
            else:
                print 'We do not have S6 parameter for the functional you use. Use 1.0 now, please correct it manually'
                dftParam='1.0 20. 25.\n'
            self.grimme='GRIMME\n'+dftParam+str(len(elements))+'\n'+elementParam        
        if args.supercell:
            self.supercell="SUPERCEL\n"+args.supercell.split(',')[0]+"\n"+args.supercell.split(',')[1]+"\n"+args.supercell.split(',')[2]+'\n'
        #render basis set
        atomList=list(set([i[0] for i in self.coord]))
        if args.basis.lower() in ['6-21gd','3-21gd','6-21g','3-21g','6-21gdp','3-21gdp']:
            self.basis='END\n'
            coreG=args.basis.lower().split('-')[0]
            polar=''
            polarH=''
            if 'd' in args.basis.lower().split('g')[1]:
                polar='2 3 1 0. 1.\n'
            if 'p' in args.basis.lower().split('g')[1]:
                polarH='2 2 1 0. 1.\n'
            for atom in atomList:
                atom=int(atom)
                if atom < 2:
                    shell=2+len(polarH.split('\n'))-1
                elif atom < 10:
                    shell=3+len(polar.split('\n'))-1
                elif atom < 18:
                    shell=4+len(polar.split('\n'))-1
                self.basis=self.basis+str(atom)+" "+str(shell)+"\n"
                if atom < 2:
                    self.basis=self.basis+"2 0 2 1. 1.\n2 0 1 0. 1.\n"
                    self.basis=self.basis+polarH       
                elif atom < 10:
                    valence=str(atom-2)
                    self.basis=self.basis+"2 0 "+coreG+" 2. 1.\n"
                    self.basis=self.basis+"2 1 2 "+valence+" 1.\n"
                    self.basis=self.basis+"2 1 1 0. 1.\n"
                    self.basis=self.basis+polar
                elif atom < 18:
                    valence=str(atom-10)
                    self.basis=self.basis+"2 0 "+coreG+" 2. 1.\n2 1 "+coreG+" 8. 1.\n"
                    self.basis=self.basis+"2 1 2 "+valence+" 1.\n"
                    self.basis=self.basis+"2 1 1 0. 1.\n"
                    self.basis=self.basis+polar
                else:
                    print "can not generate basis set for atom "+str(atom)+", please input manually"
                    self.basis=self.basis+"Input basis set for "+str(atom)+"here." 
            self.basis=self.basis+'99 0\nEND\n'
        elif "F:" in args.basis.upper():
            basis_file=open(args.basis.split(':')[1],'r').read()
            basis_list=re.split('\n\s*\n',basis_file)
            basis_string=''
            for atom in atomList:
                for basis in basis_list:
                    if str(atom)==basis.strip().split()[0].strip():
                        basis_string=basis_string+basis
            self.basis='END\n'+basis_string+'\n99 0\nEND\n'
            if not basis_string:
                print "WARNING: basis set empty!"
            #for atom in atomList:
                
        else:    
            self.basis="BASISSET\n"+args.basis+"\n"
        self.shrink="SHRINK\n"+shrink+" "+shrink+"\n"
        if args.levelshift=='0':
            self.scf="FMIXING\n"+args.fmixing+"\n"+"BROYDEN\n"+args.broyden+"\n"+"TOLINTEG\n"+args.tol+"\n"
        else:
            self.scf="LEVSHIFT\n"+args.levelshift+" 0"+"\nFMIXING\n"+args.fmixing+"\n"+"TOLINTEG\n"+args.tol+"\n"
        if args.opt != 'False':
            optKeywordDick={'A':'ATOMONLY','F':'FULLOPTG','CE':'CELLONLY','I':'ITATOCEL','CV':'CVOLOPT'}
            try:
                optKeyword=optKeywordDick[args.opt.upper()[0]]
            except KeyError:
                try:
                    optKeyword=optKeywordDick[opt[:2]]
                except KeyError:
                    print 'Can not recognize opt key word: '+opt
                    sys.exit()
            self.opt="OPTGEOM\n"+optKeyword+"\nEND\n"
        else:
            self.opt=''
        #print opt
        #print self.opt
        self.symGroup=args.sym
        self.demension=args.dim
        if external:
            self.coord=None
        if args.guess:
            self.guess='GUESSP\n'
        if args.maxcycle:
            self.cycle="MAXCYCLE\n"+args.maxcycle+"\n"


        
    def writed12(self,d12FileName):
        #for polymer or slab, change fraction coordinate to Angstrom
        
        if self.coord:
            cellParam=str(self.cell_a)
            a,b,c=1,1,1
            if self.demension=='1':
                b=float(self.cell_b)
                c=float(self.cell_c)
                self.geom="POLYMER\n"+str(self.symGroup)+"\n"+cellParam+"\n"+str(self.atoms)+"\n"
            elif self.demension=='2':
                #if self.cell_a != self.cell_b:
                cellParam=cellParam+" "+self.cell_b
                cellParam=cellParam+" "+self.cell_gamma 
                c=float(self.cell_c)
                self.geom="SLAB\n"+str(self.symGroup)+"\n"+cellParam+"\n"+str(self.atoms)+"\n"
            else:
                #if float(self.cell_a) != float(self.cell_b):
                cellParam=cellParam+" "+self.cell_b
                #if float(self.cell_a) != float(self.cell_c) and float(self.cell_b) != float(self.cell_c):  
                cellParam=cellParam+" "+self.cell_c  
                #if float(self.cell_alpha) != 90:
                cellParam=cellParam+" "+self.cell_alpha 
                #if float(self.cell_beta) != 90:
                cellParam=cellParam+" "+self.cell_beta   
                #if float(self.cell_gamma) != 90:
                cellParam=cellParam+" "+self.cell_gamma                     
                self.geom="CRYSTAL\n"+"0 0 0\n"+str(self.symGroup)+"\n"+cellParam+"\n"+str(self.atoms)+"\n"
            for coord in self.coord:
                self.geom=self.geom+"%d\t%f %f %f\n" % (coord[0],coord[1]*a,coord[2]*b,coord[3]*c)
        else:
            self.geom='EXTERNAL\n'
        d12File=open(d12FileName,'w')
        d12File.write(self.title+'\n')
        d12File.write(self.geom)
        d12File.write(self.supercell)
        d12File.write('STRUCPRT\n')
        d12File.write(self.opt)
        d12File.write(self.basis)  
        d12File.write(self.dft)
        d12File.write(self.grimme)
        d12File.write(self.shrink)
        d12File.write(self.scf)
        d12File.write(self.cycle)
        d12File.write(self.guess)
        d12File.write('END\n')   
        d12File.close()

        #for line in self.geom:
        #    print "%d\t%f %f %f\n" % (line[0],line[1],line[2],line[3]),
parser=argparse.ArgumentParser(description='Generate input file for crystal 14. Read cif file generated by Material Stuido, or use .gui file from pervious run ')
parser.add_argument("-f","--dft",metavar="DFT FUNCTIONAL", help="Set the functional, e.g. 'pbe|pbe'for X and C. 'b3lyp' for Standalone. 'b3lyp,wb97x,m062x,pbe|pbe' to benchmark different functionals. Default: m062x",  type=str,default='m062x')
parser.add_argument("-b","--basis",metavar="BASIS SET", help="Set the basis set (CUSTOM,F:basis_set_file,STO-3G,STO-6G,POB-DZVP,POB-TZVP,POB-DZVPP,6-21gdp,3-21gdp), e.g. 'STO-3G,POB-DZVP,POB-TZVP' to benchmark different functionals. Default: POB-TZVP",  type=str,default='POB-TZVP')  
parser.add_argument("-k","--shrink",help="set the shrink(k points), e.g. '2,3,4,5,6 to benchmark different K points' Default: 2",  type=str,default='2') 
parser.add_argument("-l","--levelshift",help="Set the level shift to help scf converge(1-10), set this value other than 0 will turn on broyden automatic. Default: 0", type=str,default='0') 
parser.add_argument("-c","--maxcycle",help="Set the max cycle of scf. Default:100",type=str,default='100')
parser.add_argument("-m","--fmixing",help='Set the fmixing to help scf converge(1-50)%%. Default: 50', type=str,default='50')
parser.add_argument("-d","--dim",help='Set the dimension for periodicity (1,2,3). Default: 1', type=str,default='1')
parser.add_argument("-s","--sym",help='Set the label for symmetry group. Default: 1', type=str,default='1')
parser.add_argument("-o","--opt",help='Set the paramters for geometry optimization (A:ATOMONLY,F:FULLOPTG,CE:CELLONLY,I:ITATOCEL,CV:CVOLOPT,False). Default: FULLOPTG',type=str,default='False')
parser.add_argument("-g","--guess",help='Read initial guess from previous run', action='store_true')
parser.add_argument("-D","--directory",help='Creat new directory with name filename-dirname',type=str)
parser.add_argument("--broyden",help="Using Broyden scheme to help scf converge, input 3 params seperated by space. Default: 0.0001 50 2", type=str,default='0.0001 50 2') 
parser.add_argument("--tol",help='Set the TOLINTEG value to change the truncation criteria for 2e integral, input 5 value seperated by space. Default:10 10 10 10 12', type=str,default='10 10 10 10 12')
parser.add_argument("--supercell",help='Input 3 group of number seperated by comma, e.g. 2 0 0,0 2 0,0 0 2, to creat a supercell, default: None',type=str)
parser.add_argument("--grimme",help='use grimme gd2 correction. input element number to turn on, e.g. 6,1,8 for C,H,O Or input \"A\" to determine element automatically (will fail when using external geometry), default: 0',default=None,type=str)
parser.add_argument("--freq",help='do frequency calculation. input any non-zero number to turn on, default: 0',default=0,type=int)
parser.add_argument("-v","--version",action='version', version='%(prog)s 2.0')
parser.add_argument("structureFile",nargs='*')
args=parser.parse_args()
dfts=args.dft.split(',')
basises=[]
for basis in args.basis.split(','):
    if 'F:' in basis:
        basis=os.path.basename(basis.split(':')[1]).split('.')[0]
    basises.append(basis)
shrinks=args.shrink.split(',')
dirSuffix=args.directory
guess=args.guess            
for structure in args.structureFile:
    if not os.path.isfile(structure):
        print structure+" do not exist"
        sys.exit()
    else:
        strucName=''.join(re.split('[/.]',structure)[-2:-1])
        strucSuffix=''.join(re.split('[/.]',structure)[-1])
        if strucSuffix.lower()=='gui':
            external=os.path.abspath(structure)  
            strucName='_'.join(strucName.split('_')[:-3])    
        else:
            external=''
            
    currentDIR=os.getcwd()
    if dirSuffix:
        fileDir=os.path.dirname(os.path.abspath(structure))
        
        dirName=strucName+"-"+dirSuffix
        if not os.path.exists(dirName):
            os.makedirs(dirName)
        try:
            shutil.copy(structure,dirName)
        except IOError:
            pass
        os.chdir(dirName)
    for dft in dfts:
        for basis in basises:
            for shrink in shrinks:     
                combinedName=strucName+"_"+''.join(dft.split('|'))+"_"+basis+"_"+shrink            
                inpgen=crysInpGen(combinedName)
                if not external:
                    inpgen.readcif(structure)
                if strucSuffix.lower()=='gui':
                    inpgen.readgui(external)
                inpgen.crysParamParser(args,external)
                inpgen.writed12(combinedName+'.d12')
                if external:
                    if not os.path.isfile(combinedName+'.gui'):
                        try: shutil.copy(external,combinedName+'.gui')
                        except: print(external+' not found')
                if guess:
                    if not os.path.isfile(combinedName+'.f9'):
                        if os.path.isfile(external[:-4]+'.f9'):
                            shutil.copy(external[:-4]+'.f9',combinedName+'.f9')
                        else: 
                            print(external[:-4]+'.f9 not found. Will use '+combinedName+'.f9')
    os.chdir(currentDIR)

                
